#ifndef BLOCKWISE_INTEGRATED_SOLVER_THREAD
#define BLOCKWISE_INTEGRATED_SOLVER_THREAD

#include "meshTypes_fwd.hpp"
#include "macros.hpp"
#include "bi_timing.hpp"
#include "sweSolver.hpp"
#include "bi_block.hpp"
#include "bi_edgeRow.hpp"
#include "bi_constants.hpp"
#include <experimental/memory>
#include <limits>

namespace hms::swe
{

using std::experimental::observer_ptr;
using std::experimental::make_observer;

template<typename MeshType>
class SolverThread {
public:
	using SolverType = SWESolver<MeshType>;
	static constexpr bool
		isUniMesh  { std::is_same_v<MeshType, UniMesh> },
		isRectMesh { std::is_same_v<MeshType, RectMesh> },
		isOrtho    { isUniMesh || isRectMesh };
	static constexpr int left {0}, right {1};

private:
	Array2i m_maxBlocksize;

	/* Thread-local variables. Each thread tracks the min time step for its
	 * blocks, and has storage, fluxes and possible helper vars according
	 * to the (max) block size */
	Array3Xs m_qBuf;
	Array3Xs m_fluxBuf;
	Array1Xs m_hBuf;
	ArrayX3s m_fricBuf;
	ArrayXXs m_edgesValsBuf;
	#ifdef BLOCKWISE_INTEGRATED_PRECALC_DIFFS
	ArrayXXs m_phiDiffBuf;
	#endif
	ArrayXXs m_fluxLineBuf;
	ArrayXXs m_hllcBuf;

	/* These are shared across threads. The thread is not responsible for the
	 * associated memory, but may write to it. */
	observer_ptr<const SolverType> m_solver {nullptr};
	observer_ptr<ArrayXXs> m_dNew {nullptr};
	observer_ptr<ArrayXXs> m_vNew {nullptr};

	scalar m_timestep;
	ExecTimings m_execTimings { ExecTimings::Zero() };

public:
	// Index nEdges {0};

	#ifdef BLOCKWISE_INTEGRATED_TIME_FLUXES
	FluxTimings fluxTimings { FluxTimings::Zero() };
	Timer fluxTimer;
	#define FLUX_TIMER_RESET fluxTimer.reset()
	#define TIME_FLUX(STEP) setTiming(fluxTimings, FluxStep::STEP, fluxTimer)
	#else
	#define FLUX_TIMER_RESET
	#define TIME_FLUX(STEP)
	#endif

public:
	SolverThread() = default;
	SolverThread( const IndexHelper& ih ) :
		m_qBuf         { 3, nReadCellsMax(ih) },
		m_fluxBuf      { 3, nReadCellsMax(ih) },
		m_hBuf         { 1, nReadCellsMax(ih) },
		m_fricBuf      { ih.blocks.sizeX, 3 },
		m_edgesValsBuf { ih.blocks.sizeX+1, 2*nReconstComp },
		#ifdef BLOCKWISE_INTEGRATED_PRECALC_DIFFS
		m_phiDiffBuf   { ih.blocks.sizeX + (ih.blocks.sizeX == 1), 4*3 },
		#endif
		m_fluxLineBuf  { 3*(ih.blocks.sizeX+1), 2 },
		m_hllcBuf      { ih.blocks.sizeX+1, nHllcComp }
	{
		if constexpr ( !(
			std::is_same_v<MeshType, UniMesh> ||
			std::is_same_v<MeshType, RectMesh>
		) ){
			assert(false && "SolverThread not implemented for non-orthogonal meshes!");
		}
	}
	SolverThread(
		const SolverType& solverArg,
		ArrayXXs& dNewArg,
		ArrayXXs& vNewArg,
		const IndexHelper& ih
	) :
		SolverThread(ih)
	{
		solver(solverArg);
		dNew(dNewArg);
		vNew(vNewArg);
		resetTimestep();
	}

	void solver( const SolverType& solverArg ){
		m_solver = make_observer( &solverArg );
	}
	void dNew( ArrayXXs& dNewArg ){
		assert( dNewArg.rows() == 1 );
		m_dNew = make_observer( &dNewArg );
	}
	void vNew( ArrayXXs& vNewArg ){
		assert( vNewArg.rows() == 2 );
		m_vNew = make_observer( &vNewArg );
	}

	auto solver() const -> const SolverType& {
		assert( m_solver );
		return *m_solver;
	}
	auto dNew() -> ArrayXXs& {
		assert( m_dNew );
		return *m_dNew;
	}
	auto vNew() -> ArrayXXs& {
		assert( m_vNew );
		return *m_vNew;
	}
	auto timestep() const -> scalar { return m_timestep; }
	void resetTimestep() {
		m_timestep = std::numeric_limits<scalar>::max();
	}

	auto execTimings() const -> const ExecTimings& {
		return m_execTimings;
	}
	auto execTimings() -> ExecTimings& { return m_execTimings; }

	auto qBuf   () const -> const Array3Xs& { return m_qBuf; }
	auto fluxBuf() const -> const Array3Xs& { return m_fluxBuf; }
	auto hBuf   () const -> const Array1Xs& { return m_hBuf; }

	auto qBuf   () -> Array3Xs& { return m_qBuf; }
	auto fluxBuf() -> Array3Xs& { return m_fluxBuf; }
	auto hBuf   () -> Array1Xs& { return m_hBuf; }

	template<bool singleCell=false, BlockLocation loc>
	void computeBlock( Block<loc>&& );

	template<BlockLocation loc>
	void computeStorage(
		ArrayXXsMap& qBlock,
		ArrayXXsMap& hBlock,
		const Block<loc>&,
		const Ref<const ArrayXXs>& dBlock,
		const Ref<const ArrayXXs>& zBlock,
		const Ref<const ArrayXXs>& vBlock
	);

	template<BlockLocation loc>
	scalar timestep(
		const Block<loc>& block,
		const Ref<const ArrayXXs>& dBlock,
		const Ref<const ArrayXXs>& vBlock
	);

	template<BlockLocation loc>
	void addMassSources(
		const Block<loc>& block,
		Array3XsMap& qBuf
	);

	template<BlockLocation loc>
	void computeFriction(
		const Block<loc>& block,
		// Array3XsMap& qBuf,
		Array3XsMap& fluxBuf,
		const Ref<const ArrayXXs>& dBlock,
		const Ref<const ArrayXXs>& vBlock
	);

	/* for some unknown reason, forwarding the arguments was much, much slower
	 * than typing them out. It behaved as though the Eigen::Ref arguments
	 * were copied rather than passed as references */
	// template<BlockLocation loc, typename ... Ts>
	// void computeFluxes( const Block<loc>& block, Ts&& ... args );
	template<BlockLocation loc>
	void computeFluxes(
		const Block<loc>& block,
		const Array3XsMap& q,
		      Array3XsMap& fluxes,
		const Array1XsMap& h,
		const Ref<const ArrayXXs>& d,
		const Ref<const ArrayXXs>& z,
		const Ref<const ArrayXXs>& v,
		const MeshBase<MeshType>& mesh
	);

	template<BlockLocation loc>
	void computeFluxes_cell(
		const Block<loc>& block,
		const Array3XsMap& q,
		      Array3XsMap& fluxes,
		const Array1XsMap& h,
		const Ref<const ArrayXXs>& d,
		const Ref<const ArrayXXs>& z,
		const Ref<const ArrayXXs>& v,
		const MeshBase<MeshType>& mesh
	);

	template<ReconstructionType recType, EdgeRange orien, BlockLocation loc>
	void fluxScheme(
		const Block<loc>& block,
		const Array3XsMap& q,
		      Array3XsMap& fluxes,
		const Array1XsMap& h,
		const Ref<const ArrayXXs>& d,
		const Ref<const ArrayXXs>& z,
		const Ref<const ArrayXXs>& v,
		const MeshBase<MeshType>& mesh,
		Index e,
		const Array2i& lr
	);

	template<BlockLocation loc>
	void computeFluxes_vec(
		const MeshBase<MeshType>& mesh,
		const Block<loc>& block,
		[[maybe_unused]] const Array3XsMap& q,
		[[maybe_unused]] const ArrayXXsMap& qBlock,
		      Array3XsMap& fluxes,
		const Ref<const ArrayXXs>& dBlock,
		const Ref<const ArrayXXs>& zBlock,
		const Ref<const ArrayXXs>& hBlock,
		const Ref<const ArrayXXs>& vBlock
	);

	template<ReconstructionType recType, EdgeRange orien, BlockLocation loc>
	HMS_INLINE
	EdgesValues reconstruction(
		const EdgeRow<recType,orien>&,
		const MeshBase<MeshType>& mesh,
		const Block<loc>& block,
		const ArrayXXsMap& qBlock,
		const Ref<const ArrayXXs>& dBlock,
		const Ref<const ArrayXXs>& zBlock,
		const Ref<const ArrayXXs>& hBlock,
		const Ref<const ArrayXXs>& vBlock
	);

	template<EdgeRange orien>
	HMS_INLINE
	Pair<Array3XsMap> hllcScheme(
		const EdgesValues& ev,
		Index nEdges
	);

	template<ReconstructionType recType, EdgeRange orien, BlockLocation loc>
	void fluxScheme_vec(
		EdgeRow<recType,orien>&&,
		const Array2i& lrBeg_block,
		const MeshBase<MeshType>& mesh,
		const Block<loc>& block,
		[[maybe_unused]] const Array3XsMap& q,
		[[maybe_unused]] const ArrayXXsMap& qBlock,
		      Array3XsMap& fluxes,
		const Ref<const ArrayXXs>& dBlock,
		const Ref<const ArrayXXs>& zBlock,
		const Ref<const ArrayXXs>& hBlock,
		const Ref<const ArrayXXs>& vBlock
	);

	template<BlockLocation loc>
	void computeNewState(
		const Block<loc>& block,
		Array3XsMap& qMap,
		Ref<ArrayXXs> dBlockNew,
		Ref<ArrayXXs> vBlockNew
	);

private:
	Index nReadCellsMax( const IndexHelper& ih ){
		/* the MUSCL scheme requires reading from two more cells beyond the 
		 * boundary of the block */
		return (
			Array2i{ih.blocks.sizeX, ih.blocks.sizeY} + 2*ih.mesh.nOuterLayers
		).prod();
	}

	void registerFlux(
		Array3XsMap& fluxes,
		Ref<Array3s> fluxL,
		Ref<Array3s> fluxR,
		const typename MeshType::EdgeType& edge,
		const Array2i& lr,
		const MeshBase<MeshType>& mesh
	);

	template<ReconstructionType recType, EdgeRange orien, BlockLocation loc>
	void registerFlux_vec(
		Array3XsMap& fluxes,
		Ref<Array3Xs> fluxL,
		Ref<Array3Xs> fluxR,
		const EdgeRow<recType,orien>&,
		const Array2i lrBeg_block,
		const Block<loc>& block
	);
};

} // namespace hms::swe

#include "bi_computeBlock.tpp"
#include "bi_computeStorage.tpp"
#include "bi_computeFluxes.tpp"
#include "bi_timestep.tpp"
#include "bi_sources.tpp"
#include "bi_friction.tpp"
#include "bi_newState.tpp"

#endif